import os
import cv2
import numpy as np
from tqdm import tqdm
import random

def extract_frames_from_video(video_path, output_dir, max_frames=100):
    """从视频中提取帧"""
    os.makedirs(output_dir, exist_ok=True)
    cap = cv2.VideoCapture(video_path)
    
    frame_count = 0
    saved_count = 0
    while cap.isOpened() and saved_count < max_frames:
        ret, frame = cap.read()
        if not ret:
            break
            
        # 每隔几帧保存一帧，避免过于相似
        if frame_count % 5 == 0:
            output_path = os.path.join(output_dir, f"frame_{saved_count:06d}.png")
            cv2.imwrite(output_path, frame)
            saved_count += 1
            
        frame_count += 1
    
    cap.release()
    return saved_count

def add_motion_blur(image, kernel_size=15):
    """添加运动模糊"""
    # 创建运动模糊核
    kernel = np.zeros((kernel_size, kernel_size))
    kernel[int((kernel_size-1)/2), :] = np.ones(kernel_size)
    kernel = kernel / kernel_size
    
    # 应用模糊
    blurred = cv2.filter2D(image, -1, kernel)
    return blurred

def add_complex_noise(image):
    """添加复合噪声（高斯+泊松）"""
    # 高斯噪声
    gaussian_noise = np.random.normal(0, 5, image.shape).astype(np.float32)
    noisy = image.astype(np.float32) + gaussian_noise
    noisy = np.clip(noisy, 0, 255).astype(np.uint8)
    
    # 泊松噪声
    noisy = np.random.poisson(noisy.astype(np.float32) * 0.8) / 0.8
    noisy = np.clip(noisy, 0, 255).astype(np.uint8)
    
    return noisy

def add_jpeg_compression(image, quality=30):
    """添加JPEG压缩伪影"""
    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
    result, encimg = cv2.imencode('.jpg', image, encode_param)
    decoded = cv2.imdecode(encimg, 1)
    return decoded

def simulate_led_degradation(hr_frame):
    """模拟LED显示退化流程"""
    # 1. 强运动模糊（模拟拖尾）
    blurred = add_motion_blur(hr_frame, kernel_size=15)
    
    # 2. 复合噪声（模拟低灰噪点）
    noised = add_complex_noise(blurred)
    
    # 3. JPEG压缩伪影（模拟传输损耗）
    compressed = add_jpeg_compression(noised, quality=25)
    
    # 4. 分辨率降低再恢复（模拟物理限制）
    h, w = compressed.shape[:2]
    lr = cv2.resize(compressed, (w//4, h//4), interpolation=cv2.INTER_CUBIC)
    degraded = cv2.resize(lr, (w, h), interpolation=cv2.INTER_CUBIC)
    
    return degraded

def create_degradation_dataset(input_dir, output_hr_dir, output_lr_dir):
    """创建退化数据集"""
    os.makedirs(output_hr_dir, exist_ok=True)
    os.makedirs(output_lr_dir, exist_ok=True)
    
    # 收集所有图像文件
    image_files = []
    for root, dirs, files in os.walk(input_dir):
        for file in files:
            if file.endswith(('.png', '.jpg', '.jpeg')):
                image_files.append(os.path.join(root, file))
    
    if not image_files:
        print(f"在{input_dir}中未找到图像文件")
        return
    
    print(f"找到{len(image_files)}个图像文件")
    
    for img_path in tqdm(image_files[:100], desc="Processing images"):  # 限制处理数量
        hr_image = cv2.imread(img_path)
        
        if hr_image is None:
            continue
            
        # 生成退化版本
        lr_image = simulate_led_degradation(hr_image)
        
        # 使用原始文件名保存结果
        img_file = os.path.basename(img_path)
        cv2.imwrite(os.path.join(output_hr_dir, img_file), hr_image)
        cv2.imwrite(os.path.join(output_lr_dir, img_file), lr_image)

if __name__ == "__main__":
    # 直接使用REDS train_samples中的图像创建退化数据集
    input_dir = "datasets/REDS/train_samples"
    output_base = "datasets/REDS_processed"
    
    # 创建退化数据集
    output_hr_dir = os.path.join(output_base, "HR")
    output_lr_dir = os.path.join(output_base, "LR")
    
    print("开始创建退化数据集...")
    create_degradation_dataset(input_dir, output_hr_dir, output_lr_dir)
    print("数据集创建完成！")
